
#include "glodef.h"
#include "errlog.h"
#include "vecs.h"
#include "grid3d.h"
#include "f3io.h"

#include "soah_common.h"
#include "glovar.h"

#include "ns_csplit.h"

#include <stdio.h>
#include <string.h>

double current_adj_solution_norm[];
double time_inf_adj_solution_norm[];
char adj_nrms_fname[FILENAME_MAX]="adj_norms.dat";


static void make_adv_u(void);
static void adjoint_norms(void);
static void print_inf_adj_norms(void);


/****************************************************************************/
/****************************************************************************/
/****************************************************************************/

/****************************************************************************/

int solve_adjoint_fvp(const vector v_finl, scalar q_finl, vector vtT_bc,
    void (*fill_state)(void),
    vector(*fill_adj_bc)(vector this_adj_bc),
    vector(*fill_source)(vector this_source, int *test_f))
__attribute__((warn_unused_result));

int solve_adjoint_fvp(const vector v_finl, scalar q_finl, vector vtT_bc,
    void (*fill_state)(void),
    vector(*fill_adj_bc)(vector this_adj_bc),
    vector(*fill_source)(vector this_source, int *test_f))
{
    int ret=TRUE;
    int test_src;

    int vN=get_N_vec(v);
    vector src=alloc_vec(vN);
    vector vbc=alloc_vec(vN);

    (void)vtT_bc; // to avoid compiler warning

    curr_ti=total_ti;
    move_grid();
    fill_state();

    PrintLogLine(78, '=');
    PrintLog("Time %-10f\n", (double)TIME);

    /* make final adjoint pressure compatible with v_final and u */
#if defined(CS_INITIAL_PRESSURE)
    {
        fill_source(src, &test_src);
        // for(i=0;i<3*src->N;src->X[i++]*=-1);
        scale_vec(src, -1.0);
        if(test_src)
            FinalAdjointPressure_cs(solver_data, v_finl, u, vtT_bc, src, NULL, q_finl);
        else error_msg("%s: this is not available for tested source.\n");
        make_scal_mean_0(grid, elemP, q_finl);
    }
#endif

    clean_glovecs();

    copy_vec(v_finl, v);
    copy_scal(q_finl, q);
    clean_scal(qhi);

    PRINTF("\rT = %6.3f; ", (double)(TIME));

    memset(time_inf_adj_solution_norm, 0, sizeof(time_inf_adj_solution_norm));
    adjoint_norms();

    if(dump_bin>0)
    {
        create_dir(adjoint_dir());
        save_adjoint_time(adjoint_dir());
    }
    for(; curr_ti>0;)
    {
        curr_ti--;

        double tm=MPI_Wtime();

        PrintLogLine(78, '=');
        PrintLog("Time %-10f\n", (double)TIME);

        copy_vec(u0, u1);
        copy_vec(u, u0);
        copy_scal(p0, p1);
        copy_scal(p, p0);
        fill_state();

        if(!flags.moving_grid)
        {
            copy_vec(u, adv_u);
        }
        else
        {
            move_grid();
            make_adv_u();
        }
        copy_vec(v0, v1);
        copy_vec(v, v0);
        copy_scal(q0, q1);
        copy_scal(q, q0);

        vector use_src=fill_source(src, &test_src);
        vector use_vbc=fill_adj_bc(vbc);

        /* make a timestep */
        ret=TimeStep_adjoint_cs(solver_data, v, q, v0, q0, v1, q1,
            adv_u, u, use_src, use_vbc, test_src);
        if(!ret)
            break;
        make_scal_mean_0(grid, elemP, q);

        if((dump_bin>0)&&(curr_ti%dump_bin)==0)
            save_adjoint_time(adjoint_dir());

        PRINTF("\rT = %6.3f; ", (double)(TIME));

        adjoint_norms();

        FFLUSH(stdout);

        PrintLog("CPU time for this timestep: %-6.3fsec.\n", (MPI_Wtime()-tm));
    }
    print_inf_adj_norms();
    free_vec(src);
    free_vec(vbc);
    return ret;
}


/****************************************************************************/

/****************************************************************************/

static void _state_load(void)
{
    if((dump_bin>0)&&(curr_ti%dump_bin==0))
        load_state_time(state_dir());
    else
        error_msg("%s: TODO: put error message here.\n", __func__);
}

/****************************************************************************/
/****************************************************************************/
/****************************************************************************/
/*                                                                          */
/*              analtic adjoint                                             */
/*                                                                          */
/****************************************************************************/
/****************************************************************************/

/****************************************************************************/

static void _state_analytic(void)
{
    interpolate_vector_lagr(grid, elemU, TIME, exct_vel, u);
    interpolate_scalar_lagr(grid, elemP, TIME, exct_prs, p);
    make_scal_mean_0(grid, elemP, p);
}

static vector _adj_bc_analytic(vector this_adj_bc)
{
    apply_bc_dir_vector(grid, elemU, TIME, adj_bdry_vel, this_adj_bc);
    return this_adj_bc;
}

static vector _source_analytic(vector this_source, int *test_source)
{
    test_vector_field(grid, elemU, TIME, adj_srce_trm, this_source);
    if(test_source) *test_source=FALSE;
    return this_source;
}

int solve_adjoint_analytic_fvp(const vector v_finl, scalar q_finl,
    vector vtT_bc)
{
    return solve_adjoint_fvp(v_finl, q_finl, vtT_bc,
        _state_analytic, _adj_bc_analytic, _source_analytic);
}

int solve_adjoint_analytic_load_state_fvp(const vector v_finl, scalar q_finl)
{
    return solve_adjoint_fvp(v_finl, q_finl, NULL,
        _state_load, _adj_bc_analytic, _source_analytic);
}

/****************************************************************************/
/****************************************************************************/
/****************************************************************************/
/*                                                                          */
/*              periodic adjoint                                            */
/*                                                                          */
/****************************************************************************/
/****************************************************************************/
/****************************************************************************/

static void (* const _state_periodic)(void)=_state_load;

static vector _adj_bc_periodic(vector this_adj_bc)
{
    return NULL;
}

static vector _source_periodic(vector this_source, int *test_source)
{
    if(test_source) *test_source=TRUE;
    return NULL;
}

int solve_adjoint_periodic_fvp(const vector v_finl, scalar q_finl)
{
    return solve_adjoint_fvp(v_finl, q_finl, NULL,
        _state_periodic, _adj_bc_periodic, _source_periodic);
}

/****************************************************************************/
/****************************************************************************/
/****************************************************************************/
/*                                                                          */
/*              shape adjoint                                               */
/*                                                                          */
/****************************************************************************/
/****************************************************************************/
/****************************************************************************/

static void (* const _state_shape)(void)=_state_load;

static vector _adj_bc_shape(vector this_adj_bc)
{
    return NULL;
}

static vector grad_u, curl_u;

static vector _source_shape(vector this_source, int *test_source)
{
    test_grad_e1_in_grad_e2_vec(grid, elemU, u, elemU, grad_u);
    test_curl_e1_in_curl_e2_vec(grid, elemU, u, elemU, curl_u);
    // for(int i=0;i<3*u->N;i++)
    //     this_source->X[i] = alpha*grad_u->X[i] + beta*curl_u->X[i];
    clean_vec(this_source);
    axpy_vec(this_source, alpha, grad_u);
    axpy_vec(this_source, beta, curl_u);
    if(test_source) *test_source=FALSE;
    return this_source;
}

int solve_adjoint_shape_fvp(const vector v_finl, scalar q_finl)
{
    grad_u=alloc_vec(fe_ndof(grid, elemU));
    curl_u=alloc_vec(fe_ndof(grid, elemU));
    int ret=solve_adjoint_fvp(v_finl, q_finl, NULL,
        _state_shape, _adj_bc_shape, _source_shape);
    free_vec(grad_u);
    free_vec(curl_u);
    return ret;
}

/****************************************************************************/
/****************************************************************************/
/****************************************************************************/
/*                                                                          */
/*              various utility routines                                    */
/*                                                                          */
/****************************************************************************/
/****************************************************************************/

/****************************************************************************/

static void make_adv_u(void)
{
    // vector ndisp;

    assert(elemU==fe_bub||elemU==fe_bub_vec);
    // ndisp=get_ndisp(alegrid);
    // assert(get_N_vec(ndisp)==grid->NN);

    copy_vec(u, adv_u);
    /*for(int i=0; i<grid->NN; i++)
    {
        adv_u->X[i]+=ndisp->X[i]/TStep;
        adv_u->Y[i]+=ndisp->Y[i]/TStep;
        adv_u->Z[i]+=ndisp->Z[i]/TStep;
    }*/
    // axpy_vec(adv_u, 1.0/TStep, ndisp);
}

/****************************************************************************/

static void adjoint_norms(void)
{

    interpolate_vector_lagr(grid, elemU, TIME, adj_exct_vel, ev);
    interpolate_scalar_lagr(grid, elemP, TIME, adj_exct_prs, eq);
    make_scal_mean_0(grid, elemP, eq);

    current_adj_solution_norm[NORM_DUDT]=
        norm_vec_diff_l2(grid, elemU, v, v0)/TStep;

    current_adj_solution_norm[NORM_ERR_U_L2]
        =norm_vec_diff_l2(grid, elemU, v, ev);
    current_adj_solution_norm[NORM_ERR_U_H1]
        =norm_vec_diff_h1(grid, elemU, v, ev);
    current_adj_solution_norm[NORM_ERR_P_L2]
        =norm_scal_diff_l2(grid, elemP, q, eq);
    current_adj_solution_norm[NORM_ERR_P_H1]
        =norm_scal_diff_h1(grid, elemP, q, eq);

    current_adj_solution_norm[NORM_U_L2]
        =norm_vec_diff_l2(grid, elemU, v, NULL);
    current_adj_solution_norm[NORM_P_L2]
        =norm_scal_diff_l2(grid, elemP, q, NULL);
    current_adj_solution_norm[NORM_GRAD_U_L2]=norm_grad_vec(grid, elemU, v);

    current_adj_solution_norm[NORM_EXCT_U_L2]
        =norm_vec_diff_l2(grid, elemU, ev, NULL);
    current_adj_solution_norm[NORM_EXCT_P_L2]
        =norm_scal_diff_l2(grid, elemP, eq, NULL);
    current_adj_solution_norm[NORM_EXCT_GRAD_U_L2]
        =norm_grad_vec(grid, elemU, ev);

    PrintLogNorm("v-ev", "L2", current_adj_solution_norm[NORM_ERR_U_L2]);
    PrintLogNorm("q-eq", "L2", current_adj_solution_norm[NORM_ERR_P_L2]);
    PrintLogNorm("v-ev", "H1", current_adj_solution_norm[NORM_ERR_U_H1]);
    PrintLogNorm("q-eq", "H1", current_adj_solution_norm[NORM_ERR_P_H1]);
    PrintLogNorm("ev", "L2", current_adj_solution_norm[NORM_EXCT_U_L2]);
    PrintLogNorm("eq", "L2", current_adj_solution_norm[NORM_EXCT_P_L2]);
    PrintLogNorm("grad ev", "L2",
        current_adj_solution_norm[NORM_EXCT_GRAD_U_L2]);
    PrintLogNorm("v", "L2", current_adj_solution_norm[NORM_U_L2]);
    PrintLogNorm("q", "L2", current_adj_solution_norm[NORM_P_L2]);
    PrintLogNorm("grad v", "L2", current_adj_solution_norm[NORM_GRAD_U_L2]);
    PrintLogNorm("dvdt", "L2", current_adj_solution_norm[NORM_DUDT]);

    if((total_ti-curr_ti)>=flags.time_order)
        for(int i=0; i<NORM_NUMBER; i++)
            time_inf_adj_solution_norm[i]=fmax(time_inf_adj_solution_norm[i],
            current_adj_solution_norm[i]);

    if(adj_exct_vel!=NULL)
    {
        PRINTF("||v-ev||_L2 = %10.4e",
            current_adj_solution_norm[NORM_ERR_U_L2]);
    }
    else
    {
        if(flags.stationary)
        {
            PRINTF("||dv/dt||_L2 = %10.4e",
                current_adj_solution_norm[NORM_DUDT]);
        }
        else
        {
            PRINTF("||  v ||_L2 = %10.4e",
                current_adj_solution_norm[NORM_ERR_U_L2]);
        }
    }
}

static void print_inf_adj_norms(void)
{
    PrintLogLine(1, ' ');
    PrintLogLine(78, '=');
    PrintLogNorm("v-ev", "L2Linf", time_inf_adj_solution_norm[NORM_ERR_U_L2]);
    PrintLogNorm("q-eq", "L2Linf", time_inf_adj_solution_norm[NORM_ERR_P_L2]);
    PrintLogNorm("v-ev", "H1Linf", time_inf_adj_solution_norm[NORM_ERR_U_H1]);
    PrintLogNorm("q-eq", "H1Linf", time_inf_adj_solution_norm[NORM_ERR_P_H1]);
    PrintLogNorm("ev", "L2Linf", time_inf_adj_solution_norm[NORM_EXCT_U_L2]);
    PrintLogNorm("eq", "L2Linf", time_inf_adj_solution_norm[NORM_EXCT_P_L2]);
    PrintLogNorm("grad ev", "L2Linf",
        time_inf_adj_solution_norm[NORM_EXCT_GRAD_U_L2]);
    PrintLogNorm("v", "L2Linf", time_inf_adj_solution_norm[NORM_U_L2]);
    PrintLogNorm("q", "L2Linf", time_inf_adj_solution_norm[NORM_P_L2]);
    PrintLogNorm("grad v", "L2Linf",
        time_inf_adj_solution_norm[NORM_GRAD_U_L2]);
    PrintLogNorm("dv/dt", "L2Linf", time_inf_adj_solution_norm[NORM_DUDT]);
}

void print_adj_conv_norms(void)
{
    if(mpi_my_id!=0)
        return;
    FILE *F=fopen(adj_nrms_fname, "a+t");
    if(F==NULL)
        return;
    fseek(F, 0, SEEK_SET);
    if(fgetc(F)==EOF)
    {
        fprintf(F, "%-12s    %-20s    %-14s  %-14s  %-14s  %-14s\n",
            "dT", "grid", "||v-ev||L2linf", "||q-eq||L2linf",
            "||v-ev||H1linf", "||q-eq||H1linf");
    }
    else
        fseek(F, 0, SEEK_END);
    /* fprintf(F, "%-12g    %-12g", (double) TStep, (double) GStep); */
    fprintf(F, "%-12g    %-20.20s", (double)TStep, grid->gridname);
    fprintf(F, "    %12e", time_inf_adj_solution_norm[NORM_ERR_U_L2]);
    fprintf(F, "    %12e", time_inf_adj_solution_norm[NORM_ERR_P_L2]);
    fprintf(F, "    %12e", time_inf_adj_solution_norm[NORM_ERR_U_H1]);
    fprintf(F, "    %12e", time_inf_adj_solution_norm[NORM_ERR_P_H1]);
    fprintf(F, "\n");
    fclose(F);
}

